"""
聚合操作符实现
"""

from typing import Optional, List, Dict, Any, Iterator, Tuple
from collections import defaultdict
import time

from ..executor_base import IteratorExecutor
from ..types import Record, RecordBatch, ExecutionContext, AggregateFunction

class AggregateOperatorBase(IteratorExecutor):
    """聚合操作符基类"""
    
    def __init__(self, executor_id: str, context: ExecutionContext,
                 child: IteratorExecutor,
                 group_by_columns: List[str],
                 aggregate_specs: List[Dict[str, Any]]):
        super().__init__(executor_id, context)
        self.child = child
        self.group_by_columns = group_by_columns
        self.aggregate_specs = aggregate_specs
        
        # 添加子执行器
        self.add_child(child)
        
        # 验证聚合规格
        self._validate_aggregate_specs()
    
    def _validate_aggregate_specs(self):
        """验证聚合规格"""
        for spec in self.aggregate_specs:
            if 'function' not in spec:
                raise ValueError("Aggregate spec must contain 'function'")
            
            func_name = spec['function'].upper()
            if func_name not in [f.value.upper() for f in AggregateFunction]:
                raise ValueError(f"Unsupported aggregate function: {func_name}")
            
            if 'column' not in spec and func_name != 'COUNT':
                raise ValueError(f"Aggregate function {func_name} requires 'column'")
    
    def _create_group_key(self, record: Record) -> Tuple:
        """创建分组键"""
        if not self.group_by_columns:
            return ()  # 无分组，所有记录在一组
        
        key_values = []
        for col in self.group_by_columns:
            key_values.append(record.get(col))
        return tuple(key_values)
    
    def _initialize_aggregators(self) -> Dict[str, Any]:
        """初始化聚合器"""
        aggregators = {}
        
        for spec in self.aggregate_specs:
            alias = spec.get('alias', f"{spec['function']}_{spec.get('column', '*')}")
            func_name = spec['function'].upper()
            
            if func_name == 'COUNT':
                aggregators[alias] = {'type': 'count', 'value': 0}
            elif func_name == 'SUM':
                aggregators[alias] = {'type': 'sum', 'value': 0}
            elif func_name == 'AVG':
                aggregators[alias] = {'type': 'avg', 'sum': 0, 'count': 0}
            elif func_name == 'MIN':
                aggregators[alias] = {'type': 'min', 'value': None}
            elif func_name == 'MAX':
                aggregators[alias] = {'type': 'max', 'value': None}
            elif func_name == 'GROUP_CONCAT':
                aggregators[alias] = {'type': 'group_concat', 'values': []}
        
        return aggregators
    
    def _update_aggregators(self, aggregators: Dict[str, Any], record: Record):
        """更新聚合器"""
        for spec in self.aggregate_specs:
            alias = spec.get('alias', f"{spec['function']}_{spec.get('column', '*')}")
            func_name = spec['function'].upper()
            column = spec.get('column')
            
            agg = aggregators[alias]
            
            if func_name == 'COUNT':
                if column == '*' or column is None:
                    agg['value'] += 1
                elif record.get(column) is not None:
                    agg['value'] += 1
            
            elif func_name in ['SUM', 'AVG']:
                value = record.get(column)
                if value is not None:
                    try:
                        numeric_value = float(value)
                        if func_name == 'SUM':
                            agg['value'] += numeric_value
                        else:  # AVG
                            agg['sum'] += numeric_value
                            agg['count'] += 1
                    except (ValueError, TypeError):
                        pass
            
            elif func_name == 'MIN':
                value = record.get(column)
                if value is not None:
                    if agg['value'] is None or value < agg['value']:
                        agg['value'] = value
            
            elif func_name == 'MAX':
                value = record.get(column)
                if value is not None:
                    if agg['value'] is None or value > agg['value']:
                        agg['value'] = value
            
            elif func_name == 'GROUP_CONCAT':
                value = record.get(column)
                if value is not None:
                    agg['values'].append(str(value))
    
    def _finalize_aggregators(self, aggregators: Dict[str, Any]) -> Record:
        """完成聚合计算"""
        result = {}
        
        for alias, agg in aggregators.items():
            agg_type = agg['type']
            
            if agg_type == 'count':
                result[alias] = agg['value']
            elif agg_type == 'sum':
                result[alias] = agg['value']
            elif agg_type == 'avg':
                if agg['count'] > 0:
                    result[alias] = agg['sum'] / agg['count']
                else:
                    result[alias] = None
            elif agg_type in ['min', 'max']:
                result[alias] = agg['value']
            elif agg_type == 'group_concat':
                separator = ','  # 默认分隔符
                result[alias] = separator.join(agg['values'])
        
        return result

class HashAggregateOperator(AggregateOperatorBase):
    """哈希聚合操作符"""
    
    def __init__(self, executor_id: str, context: ExecutionContext,
                 child: IteratorExecutor,
                 group_by_columns: List[str],
                 aggregate_specs: List[Dict[str, Any]]):
        super().__init__(executor_id, context, child, group_by_columns, aggregate_specs)
        self.hash_table: Dict[Tuple, Dict[str, Any]] = {}
    
    def open(self):
        """打开哈希聚合操作符"""
        with self.lock:
            if self.is_open:
                return
            
            self.is_open = True
            
            # 打开子执行器
            self.child.open()
            
            # 构建哈希聚合表
            self._build_hash_table()
            
            # 创建结果迭代器
            self.current_iterator = self._create_result_iterator()
    
    def _build_hash_table(self):
        """构建哈希聚合表"""
        while True:
            child_batch = self.child.next_batch()
            if not child_batch:
                break
            
            for record in child_batch:
                self.check_timeout()
                
                # 更新统计
                self.update_stats(rows_processed=1)
                
                # 创建分组键
                group_key = self._create_group_key(record)
                
                # 初始化或获取聚合器
                if group_key not in self.hash_table:
                    self.hash_table[group_key] = self._initialize_aggregators()
                
                # 更新聚合器
                self._update_aggregators(self.hash_table[group_key], record)
    
    def _create_result_iterator(self) -> Iterator[Record]:
        """创建结果迭代器"""
        for group_key, aggregators in self.hash_table.items():
            self.check_timeout()
            
            # 创建结果记录
            result = {}
            
            # 添加分组列
            for i, col in enumerate(self.group_by_columns):
                if i < len(group_key):
                    result[col] = group_key[i]
            
            # 添加聚合结果
            agg_result = self._finalize_aggregators(aggregators)
            result.update(agg_result)
            
            self.update_stats(rows_returned=1)
            yield result
    
    def close(self):
        """关闭哈希聚合操作符"""
        with self.lock:
            if not self.is_open:
                return
            
            # 关闭子执行器
            self.child.close()
            
            # 清理哈希表
            self.hash_table.clear()
            
            self.is_open = False
            self.current_iterator = None

class SortAggregateOperator(AggregateOperatorBase):
    """排序聚合操作符"""
    
    def __init__(self, executor_id: str, context: ExecutionContext,
                 child: IteratorExecutor,
                 group_by_columns: List[str],
                 aggregate_specs: List[Dict[str, Any]]):
        super().__init__(executor_id, context, child, group_by_columns, aggregate_specs)
        self.sorted_records: List[Record] = []
        self.current_index = 0
    
    def open(self):
        """打开排序聚合操作符"""
        with self.lock:
            if self.is_open:
                return
            
            self.is_open = True
            
            # 打开子执行器
            self.child.open()
            
            # 加载并排序所有记录
            self._load_and_sort_records()
            
            # 创建聚合迭代器
            self.current_iterator = self._create_aggregate_iterator()
    
    def _load_and_sort_records(self):
        """加载并排序记录"""
        # 加载所有记录
        while True:
            child_batch = self.child.next_batch()
            if not child_batch:
                break
            
            self.sorted_records.extend(child_batch)
            self.update_stats(rows_processed=len(child_batch))
        
        # 按分组列排序
        if self.group_by_columns:
            self.sorted_records.sort(
                key=lambda r: self._create_group_key(r)
            )
    
    def _create_aggregate_iterator(self) -> Iterator[Record]:
        """创建聚合迭代器"""
        if not self.sorted_records:
            return
        
        current_group_key = None
        current_aggregators = None
        group_records = []
        
        for record in self.sorted_records:
            self.check_timeout()
            
            group_key = self._create_group_key(record)
            
            # 检查是否开始新的分组
            if current_group_key != group_key:
                # 输出上一个分组的结果
                if current_group_key is not None:
                    result = self._create_group_result(current_group_key, current_aggregators)
                    self.update_stats(rows_returned=1)
                    yield result
                
                # 开始新分组
                current_group_key = group_key
                current_aggregators = self._initialize_aggregators()
                group_records = []
            
            # 更新当前分组的聚合器
            self._update_aggregators(current_aggregators, record)
            group_records.append(record)
        
        # 输出最后一个分组的结果
        if current_group_key is not None:
            result = self._create_group_result(current_group_key, current_aggregators)
            self.update_stats(rows_returned=1)
            yield result
    
    def _create_group_result(self, group_key: Tuple, aggregators: Dict[str, Any]) -> Record:
        """创建分组结果"""
        result = {}
        
        # 添加分组列
        for i, col in enumerate(self.group_by_columns):
            if i < len(group_key):
                result[col] = group_key[i]
        
        # 添加聚合结果
        agg_result = self._finalize_aggregators(aggregators)
        result.update(agg_result)
        
        return result
    
    def close(self):
        """关闭排序聚合操作符"""
        with self.lock:
            if not self.is_open:
                return
            
            # 关闭子执行器
            self.child.close()
            
            # 清理排序记录
            self.sorted_records.clear()
            
            self.is_open = False
            self.current_iterator = None

class StreamingAggregateOperator(AggregateOperatorBase):
    """流式聚合操作符，适用于预排序的输入"""
    
    def __init__(self, executor_id: str, context: ExecutionContext,
                 child: IteratorExecutor,
                 group_by_columns: List[str],
                 aggregate_specs: List[Dict[str, Any]]):
        super().__init__(executor_id, context, child, group_by_columns, aggregate_specs)
        self.current_group_key = None
        self.current_aggregators = None
        self.pending_record = None
    
    def open(self):
        """打开流式聚合操作符"""
        with self.lock:
            if self.is_open:
                return
            
            self.is_open = True
            
            # 打开子执行器
            self.child.open()
            
            # 创建流式聚合迭代器
            self.current_iterator = self._create_streaming_iterator()
    
    def _create_streaming_iterator(self) -> Iterator[Record]:
        """创建流式聚合迭代器"""
        while True:
            # 获取下一条记录
            if self.pending_record is not None:
                record = self.pending_record
                self.pending_record = None
            else:
                record = self._get_next_record()
                if record is None:
                    break
            
            self.check_timeout()
            self.update_stats(rows_processed=1)
            
            group_key = self._create_group_key(record)
            
            # 检查是否开始新的分组
            if self.current_group_key != group_key:
                # 输出上一个分组的结果
                if self.current_group_key is not None:
                    result = self._create_group_result(self.current_group_key, self.current_aggregators)
                    
                    # 保存当前记录作为下一个分组的开始
                    self.pending_record = record
                    self.current_group_key = group_key
                    self.current_aggregators = self._initialize_aggregators()
                    
                    self.update_stats(rows_returned=1)
                    yield result
                    continue
                else:
                    # 第一个分组
                    self.current_group_key = group_key
                    self.current_aggregators = self._initialize_aggregators()
            
            # 更新当前分组的聚合器
            self._update_aggregators(self.current_aggregators, record)
        
        # 输出最后一个分组的结果
        if self.current_group_key is not None:
            result = self._create_group_result(self.current_group_key, self.current_aggregators)
            self.update_stats(rows_returned=1)
            yield result
    
    def _get_next_record(self) -> Optional[Record]:
        """获取下一条记录"""
        batch = self.child.next_batch()
        if batch:
            return batch[0] if batch else None
        return None
    
    def _create_group_result(self, group_key: Tuple, aggregators: Dict[str, Any]) -> Record:
        """创建分组结果"""
        result = {}
        
        # 添加分组列
        for i, col in enumerate(self.group_by_columns):
            if i < len(group_key):
                result[col] = group_key[i]
        
        # 添加聚合结果
        agg_result = self._finalize_aggregators(aggregators)
        result.update(agg_result)
        
        return result
    
    def close(self):
        """关闭流式聚合操作符"""
        with self.lock:
            if not self.is_open:
                return
            
            # 关闭子执行器
            self.child.close()
            
            # 重置状态
            self.current_group_key = None
            self.current_aggregators = None
            self.pending_record = None
            
            self.is_open = False
            self.current_iterator = None


